#include "RHICommandBuffer.h"
#include "../RenderState.h"
#include "../../Utils/LogUtil.h"
#include "../../Game/Process/RenderProcess.h"

RHICommandBuffer::BuilderDetails::BuilderDetails()
{
	queueFlag = VK_QUEUE_COMPUTE_BIT;

	commandBufferAllocateInfo = {};
	commandBufferAllocateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
	commandBufferAllocateInfo.commandPool = GRenderState->GetCommandPool(queueFlag);
	commandBufferAllocateInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
	commandBufferAllocateInfo.commandBufferCount = GRenderState->GetFrameCount();

	commandBufferUsage = VK_COMMAND_BUFFER_USAGE_SIMULTANEOUS_USE_BIT;
}

RHICommandBuffer::Builder::Builder() noexcept = default;
RHICommandBuffer::Builder::~Builder() noexcept = default;
RHICommandBuffer::Builder::Builder(RHICommandBuffer::Builder const& rhs) noexcept = default;
RHICommandBuffer::Builder::Builder(RHICommandBuffer::Builder && rhs) noexcept = default;
RHICommandBuffer::Builder& RHICommandBuffer::Builder::operator=(RHICommandBuffer::Builder const& rhs) noexcept = default;
RHICommandBuffer::Builder& RHICommandBuffer::Builder::operator=(RHICommandBuffer::Builder && rhs) noexcept = default;


RHICommandBuffer::Builder& RHICommandBuffer::Builder::SetQueue(VkQueueFlagBits queueFlag) noexcept
{
	mImpl->commandBufferAllocateInfo.commandPool = GRenderState->GetCommandPool(queueFlag);
	mImpl->queueFlag = queueFlag;
	return *this;
}

RHICommandBuffer::Builder& RHICommandBuffer::Builder::SetLevel(VkCommandBufferLevel level) noexcept
{
	mImpl->commandBufferAllocateInfo.level = level;
	return *this;
}

RHICommandBuffer::Builder& RHICommandBuffer::Builder::SetCount(uint32 count) noexcept
{
	mImpl->commandBufferAllocateInfo.commandBufferCount = count;
	return *this;
}

RHICommandBuffer::Builder& RHICommandBuffer::Builder::SetCommandBufferUsage(VkCommandBufferUsageFlagBits commandBufferUsage) noexcept
{
	mImpl->commandBufferUsage = commandBufferUsage;
	return *this;
}

VkCommandBuffer RHICommandBuffer::GetCommandBuffer(uint32 index) const
{
	if (index == UINT32_MAX)
	{
		assert(mCommandBuffers.size() == GRenderState->GetFrameCount() || mCommandBuffers.size() == 1);
		return mCommandBuffers.size() == 1 ? mCommandBuffers[0] : mCommandBuffers[GRenderState->GetFrameIndex()];
	}
	assert(index < mCommandBuffers.size());
	return mCommandBuffers[index];
}

VkQueueFlagBits RHICommandBuffer::GetQueueFlag() const
{
	return mQueueFlag;
}

VkCommandBuffer RHICommandBuffer::BeginCommandBuffer(uint32 index) const
{
	VkCommandBuffer commandBuffer = GetCommandBuffer(index);

	VK_CHECK_RESULT(vkBeginCommandBuffer(commandBuffer, &mCommandBufferBeginInfo));

	return commandBuffer;
}

void RHICommandBuffer::FlushCommandBuffer(VkCommandBuffer commandBuffer, std::vector<RefCountPtr<RHISemaphore>> waitSemaphores, std::vector<RefCountPtr<RHISemaphore>> signalSemaphores, RefCountPtr<RHIFence> fence) const
{
	// Make sure the command buffer is in this rhi resource
	bool isExist = false;
	for (uint32 i = 0; i < mCommandBuffers.size(); ++i)
	{
		if (commandBuffer == mCommandBuffers[i])
		{
			isExist = true;
			break;
		}
	}
	assert(isExist);

	VK_CHECK_RESULT(vkEndCommandBuffer(commandBuffer));


	std::vector<VkSemaphore> waitVkSemaphores;
	std::vector<VkPipelineStageFlags> waitDstStages;

	waitVkSemaphores.resize(waitSemaphores.size(), VK_NULL_HANDLE);
	waitDstStages.resize(waitSemaphores.size());

	for (uint32 i = 0; i < waitSemaphores.size(); ++i)
	{
		waitVkSemaphores[i] = waitSemaphores[i]->GetSemaphore();
		waitDstStages[i] = waitSemaphores[i]->GetWaitDstStage();
	}

	std::vector<VkSemaphore> signalVkSemaphores;
	signalVkSemaphores.resize(signalSemaphores.size(), VK_NULL_HANDLE);
	for (uint32 i = 0; i < signalSemaphores.size(); ++i)
	{
		signalVkSemaphores[i] = signalSemaphores[i]->GetSemaphore();
	}

	VkSubmitInfo submitInfo{};
	submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;

	submitInfo.waitSemaphoreCount = static_cast<uint32>(waitVkSemaphores.size());
	submitInfo.pWaitSemaphores = waitVkSemaphores.data();

	submitInfo.pWaitDstStageMask = waitDstStages.data();

	submitInfo.commandBufferCount = 1;
	submitInfo.pCommandBuffers = &commandBuffer;

	submitInfo.signalSemaphoreCount = static_cast<uint32>(signalVkSemaphores.size());
	submitInfo.pSignalSemaphores = signalVkSemaphores.data();

	VK_CHECK_RESULT(vkQueueSubmit(GVulkanRHI->GetQueue(mQueueFlag), 1, &submitInfo, fence.IsValid() ? fence->GetFence() : VK_NULL_HANDLE));
}

RefCountPtr<RHICommandBuffer> RHICommandBuffer::Builder::Build()
{
	assert(IsInRenderThread());
	return RefCountPtr<RHICommandBuffer>(NewObject<RHICommandBuffer>(*mImpl));
}

RefCountPtr<RHICommandBuffer> RHICommandBuffer::Builder::Build(RenderProcess * process, const std::string & name)
{
	assert(IsInRenderThread());
	RHICommandBuffer* commandBuffer = NewObject<RHICommandBuffer>(*mImpl);
	process->AddCommandBuffer(name, commandBuffer);
	return RefCountPtr<RHICommandBuffer>(commandBuffer);
}

void RHICommandBuffer::ReleaseResource()
{
	if (IsInRenderThread())
	{
		if (mProcess != nullptr)
		{
			mProcess->RemoveCommandBuffer(mName, this);
		}
		DeleteObject<RHICommandBuffer>(this);
	}
	else
	{
		AsyncTask(GRenderThreadId, [this]() {
			if (mProcess != nullptr)
			{
				mProcess->RemoveCommandBuffer(mName, this);
			}
			DeleteObject<RHICommandBuffer>(this);
			});
	}
}

RHICommandBuffer::~RHICommandBuffer()
{
	GRenderState->DelayGCCommandBuffer(mQueueFlag, mCommandBuffers);
}

RHICommandBuffer::RHICommandBuffer(BuilderDetails & builderDetails)
{
	mQueueFlag = builderDetails.queueFlag;

	mCommandBufferAllocateInfo = builderDetails.commandBufferAllocateInfo;

	mCommandBuffers.resize(mCommandBufferAllocateInfo.commandBufferCount, VK_NULL_HANDLE);

	VK_CHECK_RESULT(vkAllocateCommandBuffers(GVulkanRHI->GetLogicalDevice(), &mCommandBufferAllocateInfo, mCommandBuffers.data()));

	mCommandBufferBeginInfo = {};
	mCommandBufferBeginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
	mCommandBufferBeginInfo.flags = builderDetails.commandBufferUsage;
}